# 对位姿精度进行验证 
# 1. 生成标定板的位姿

import numpy as np
import os,sys
from pose_evaluator_lmo import PoseEvaluatorLMO

def read_april_tag_pose(pose_dir,iter):
    """
    读取AprilTag的位姿
    :return: 位姿列表
    """
    pose = np.loadtxt(f'{pose_dir}/pose_april/{str(iter)}.txt')
    pose = pose.reshape(4, 4)
    return pose[:3,:]

def read_halcon_pose(pose_dir,iter):
    """
    读取halcon的位姿
    :return: 位姿列表
    """
    def process_pose(pose):
        """
        处理位姿, 转换位姿
        :return: 旋转矩阵的逆
        """
        pose = np.linalg.inv(pose)
        return pose
    pose = np.loadtxt(f'{pose_dir}/pose_halcon/{str(iter)}.txt')
    pose = pose.reshape(3, 4)
    # pose = np.vstack((pose, np.array([0, 0, 0, 1])))
    # pose = process_pose(pose)
    return pose


def error_rotation(pose1, pose2):
    """
    计算旋转误差
    :return: 旋转误差
    """
    R1 = pose1[:3, :3]
    R2 = pose2[:3, :3]
    return np.arccos((np.trace(np.dot(R1, R2.T)) - 1) / 2)


def get_pose_info(pose_dir):
    """
    读取AprilTag的位姿
    :return: 位姿字典
    """
    assert os.path.exists(pose_dir), "The pose dir does not exist"
    assert len(os.listdir(f'{pose_dir}\pose_april')) == len(os.listdir(f'{pose_dir}\pose_halcon')),"The number of pose_april and pose_halcon is not equal"
    
    len_pose = len(os.listdir(f'{pose_dir}\pose_april'))

    pose_aprils = []
    pose_halcons = []
    
    for i in range(len_pose):
        pose_april = read_april_tag_pose(pose_dir,i)
        pose_halcon = read_halcon_pose(pose_dir,i)
        # pose_halcon = process_pose(pose_halcon)
        pose_aprils.append(pose_april)
        pose_halcons.append(pose_halcon)
    return {1:pose_aprils},{1:pose_halcons} # 返回字典


def get_april_tag_info(dataset_dir):

    def load_model_info(points):
        """
        Load information about the 3D model from the BOP files
        """
        infos = {}
        extents = 2 * np.max(np.absolute(points), axis=0)
        infos['diameter'] = np.sqrt(np.sum(extents * extents))
        infos['min_x'], infos['min_y'], infos['min_z'] = np.min(points, axis=0)
        infos['max_x'], infos['max_y'], infos['max_z'] = np.min(points, axis=0)
        return infos
    
    models_pcd = np.load(os.path.join(dataset_dir, "models_pcd.npy"))
    models = {1: {"pts":models_pcd[0]}}#[n,1024,3]
    models_info={1:load_model_info(models_pcd[0])}
    models_symmetry={1:1} # if the object is symmetric == 1 else 0
    classes = [1,]

    return {
        "models": models,
        "classes": classes,
        "models_info": models_info,
        "models_symmetry": models_symmetry,
    }


def main():
    pose_dir = r'E:\pose\datasets\obj_ac\obj_000000'
    april_model_dir =  r'E:\pose\datasets\obj_ac\obj_000000\model'

    info = get_april_tag_info(april_model_dir)
    metrics = PoseEvaluatorLMO(models=info['models'], \
                                classes=info['classes'],\
                                model_info=info['models_info'], \
                                model_symmetry=info['models_symmetry'],\
                                depth_scale=1)
    
    gt_pose,halcon_pose=get_pose_info(pose_dir)

    metrics_save_dir ='E:/tools/eval/results/'
    metrics.reset()
    metrics.poses_gt ={1:gt_pose[1]}
    metrics.poses_pred = {1:halcon_pose[1]}
    metrics.evaluate_pose_adds(metrics_save_dir)
    metrics.calculate_class_avg_rotation_error(metrics_save_dir)
    metrics.calculate_class_avg_translation_error(metrics_save_dir)
    metrics.calc_mAP_r(metrics_save_dir)
    metrics.calc_mAP_t(metrics_save_dir)
    print('Done')


if __name__ == '__main__':
    main()

    